type side = BUY | SELL

type outright_id = OUT1 | OUT2 | OUT3
type strategy_id = STRAT1 | STRAT2
type month = Mar | Jun | Sep | Dec

(* Map an outright with an expiry *)
let contract_expiry = function
  | OUT1 -> Mar
  | OUT2 -> Jun
  | OUT3 -> Sep

(* Convert month to an integer *)
let month_to_int = function
  | Mar -> 3
  | Jun -> 6
  | Sep -> 9
  | Dec -> 12
;;

(* Return true of m1 is nearer (or equal) to m2 *)
let month_comp (m1 : month) (m2 : month) =
  (month_to_int m1) < (month_to_int  m2)
;;

type instrument =
  | Strategy of strategy_id
  | Outright of outright_id

(* Level information *)
type level_info = {
  li_qty : int
  ; li_price : int
}

(* best bid and ask information *)
type best_bid_ask = {
  bid_info : level_info option
  ; ask_info : level_info option
}

(* Best bid/ask for all of the books *)
type books_info = {
  book1 : best_bid_ask
  ; book2 : best_bid_ask
  ; book3 : best_bid_ask
}

(* Order type *)
type order = {
  o_qty : int
  ; o_price : int
  ; o_time : int
  ; o_id : int
  ; o_side : side
  ; o_client_id : int
  ; o_inst : instrument
  ; o_is_implied : bool
}

(* Helper function to make order creation simpler *)
let make si qty price id inst clientid isimp time =
  {o_qty = qty ; o_price = price; o_id = id; o_side = si;
   o_client_id = clientid; o_inst = inst; o_is_implied = isimp;
   o_time = time }


(* outright order book *)
type book = {
  b_buys : order list
  ; b_sells : order list
}

let empty_book = { b_buys = []; b_sells = [] }

(* Individual leg *)
type leg = {
  leg_sec_idx : outright_id
  ; leg_mult : int
}

(* Strategy is composed of legs *)
type strategy = {
  time_created : int
  ; leg1 : leg
  ; leg2 : leg
  ; leg3 : leg
}

(* Helper function to make strategy creation smaller *)
let make_strat tcreated m1 m2 m3 = {
  time_created = tcreated
  ; leg1 = { leg_sec_idx = OUT1; leg_mult = m1 }
  ; leg2 = { leg_sec_idx = OUT2; leg_mult = m2 }
  ; leg3 = { leg_sec_idx = OUT3; leg_mult = m3 }
}

type implied_strat_ord = {
  max_strat : int option
  ; strat_price : int option
}

(* New order message *)
type new_ord_msg = {
  no_client_id : int
  ; no_inst_type : instrument
  ; no_qty : int
  ; no_side : side
  ; no_price : int
}

(* cancel order ID *)
type cancel_ord_msg = {
  co_client_id : int
  ; co_order_id : int
  ; co_instrument : instrument
  ; co_side : side
}

(* Inbound messages type *)
type inbound_msg =
  | NewOrder of new_ord_msg
  | CancelOrder of cancel_ord_msg
  | ImpliedUncross

(* Helper function for creating new order messages *)
let make_no_msg cid inst qty sd p =
 NewOrder {
  no_client_id = cid
  ; no_inst_type = inst
  ; no_qty = qty
  ; no_side = sd
  ; no_price = p
 }

(* ack message *)
type ack_msg = {
  ack_client_id : int
  ; ack_order_id : int
  ; ack_inst_type : instrument
  ; ack_qty : int
  ; ack_side : side
  ; ack_price : int
}

(* fill information *)
type fill = {
  fill_client_id : int
  ; fill_qty : int
  ; fill_price : int
  ; fill_order_id : int
  ; fill_order_done : bool
}

(* uncross result *)
type uncross_res = {
  uncrossed_book : book
  ; uncrossed_fills : fill list
  ; uncrossed_qty : int
}

(* outbound message type *)
type outbound_msg =
  | Ack of ack_msg
  | Fill of fill
  | UncrossResult of uncross_res
;;

(* The entire market - strategy definitions, order books, messages, etc. s*)
type market = {

  (* current time*)
  curr_time : int

  (* used for order ID counter *)
  ; last_ord_id : int

  (* two strategy definitions *)
  ; strat1    : strategy
  ; strat2    : strategy

  (* outright books *)
  ; out_book1 : book
  ; out_book2 : book
  ; out_book3 : book

  (* strategy books *)
  ; s_book1   : book
  ; s_book2   : book

  (* inbound and outbound message queues *)
  ; inbound_msgs  : inbound_msg list
  ; outbound_msgs : outbound_msg list

}

(* Here's an example of a custom printer that we can install for arbitrary data types. *)

#program;;
#require "tyxml";;
let html_of_order (o : order) =
  let module H = Tyxml.Html in
  H.div
  ~a:(if o.o_is_implied then [H.a_style "color: red"] else [])
  [ H.div
    ~a:[H.a_style "font-size: 1.4em"]
    [H.txt (Format.asprintf "%s (%s)" (Z.to_string o.o_price) (Z.to_string o.o_qty))]
  ; H.div (if o.o_is_implied then [H.txt "Implied"] else [])
  ]

let html elt =
  let module H = Tyxml.Html in
  Document.html (Document.Unsafe_.html_of_string @@ CCFormat.sprintf "%a" (H.pp_elt ()) elt);;

let doc_of_order (o:order) =
  let module H = Tyxml.Html in
  html (H.div [html_of_order o]);;

#install_doc doc_of_order;;

let html_of_book ?(title="") (b: book) =
  let module H = Tyxml.Html in
  let rec build_rows acc buys sells =
      match buys, sells with
      | b :: bs, s :: ss -> build_rows (acc @ [H.tr [H.td [html_of_order b]; H.td [html_of_order s]]]) bs ss
      | b :: bs, [] -> build_rows (acc @ [H.tr [H.td [html_of_order b]; H.td [H.txt "-"]]]) bs []
      | [], s :: ss -> build_rows (acc @ [H.tr [H.td [H.txt "-"]; H.td [html_of_order s]]]) [] ss
      | [], [] -> acc
  in
  H.div
  ~a:[H.a_style "margin-right:1em; display: flex; flex-direction: column; align-items: center; justify-content: flex-start"]
  [ H.div ~a:[H.a_style "font-weight: bold"] [H.txt title]
  ; H.table
    ~thead:(H.thead [H.tr [H.th [H.txt "Buys"]; H.th [H.txt "Sells"]]])
    (build_rows [] b.b_buys b.b_sells)]

let doc_of_book (b:book) =
  let module H = Tyxml.Html in
  html (H.div [html_of_book ~title:"M1 Mar21" b]);;

#install_doc doc_of_book;;

let html_of_market (m: market) =
  let module H = Tyxml.Html in
  H.div
  [ H.div ~a:[H.a_style "display: flex"]
    [ html_of_book ~title:"Strategy 1" m.s_book1
    ; html_of_book ~title:"Strategy 2" m.s_book2
    ]
  ; H.div ~a:[H.a_style "margin-top: 1em; display: flex"]
    [ html_of_book ~title:"Book 1" m.out_book1
    ; html_of_book ~title:"Book 2" m.out_book2
    ; html_of_book ~title:"Book 3" m.out_book3
    ]]

let doc_of_market (m : market) =
  html (html_of_market m);;

#install_doc doc_of_market;;

#logic;;

let leg = { leg_sec_idx = OUT1; leg_mult = 1 } in

let strat = { time_created = 0; leg1 = leg; leg2 = leg; leg3 = leg } in

let b1 = {
  b_buys = [
    (make BUY 100 54 1 (Outright OUT1) 1 true 1)
    ;(make BUY 100 54 2 (Outright OUT1) 1 false 1)
  ]
  ; b_sells = [
    (make SELL 100 54 3 (Outright OUT1) 1 false 1)
    ;(make SELL 100 54 4 (Outright OUT1) 1 false 1)
  ] } in

  { curr_time = 1
  ; last_ord_id = 1
  ; strat1 = strat
  ; strat2 = strat
  ; out_book1 = b1
  ; out_book2 = b1
  ; out_book3 = b1
  ; s_book1 = b1
  ; s_book2 = b1
  ; inbound_msgs = []
  ; outbound_msgs = []
  }

(* Convert fills into outbound messages *)
let rec create_fill_msgs (f : fill list) =
  match f with
  | [] -> []
  | x::xs -> (Fill x) :: create_fill_msgs xs

(* TODO: recode this with higher-order functions *)
let rec cancel_ord_side (orders : order list) (c : cancel_ord_msg) =
  match orders with
  | [] -> []
  | x::xs ->
    begin
      if (x.o_client_id = c.co_client_id) && (x.o_id = c.co_order_id) then xs
      else x :: (cancel_ord_side xs c)
    end

(* Helper to cancel orders *)
let cancel_ord_book (co : cancel_ord_msg) (b : book) =
  match co.co_side with
  | BUY -> { b with b_buys = (cancel_ord_side b.b_buys co) }
  | SELL -> { b with b_sells = (cancel_ord_side b.b_sells co) }

(* function used insert individual orders *)
let rec insert_order_side (orders : order list) (o : order) =
  match orders with
  | [] -> [ o ]
  | x::xs ->
    begin
      if o.o_side = BUY then
        (if o.o_price > x.o_price then o :: orders else x :: (insert_order_side xs o))
      else
        (if o.o_price < x.o_price then o :: orders else x :: (insert_order_side xs o))
    end

(* insert order into the book *)
let insert_order (o : order) (b : book) =
  if o.o_side = BUY then
    { b with b_buys = (insert_order_side b.b_buys o) }
  else
    { b with b_sells = (insert_order_side b.b_sells o) }

(* The fills are adjusted to a single fill price during the uncross *)
let rec adjust_fill_prices (fills : fill list) ( f_price : int ) =
  match fills with
  | [] -> []
  | x::xs -> { x with fill_price = f_price } :: ( adjust_fill_prices xs f_price )
;;

(* Measure for proving termination of `uncross_book` below *)
let book_measure b =
  Ordinal.of_int (List.length b.b_buys + List.length b.b_sells)

let rec uncross_book (b : book) (fills : fill list) (filled_qty : int) =
  match b.b_buys, b.b_sells with
  | [], [] | _, [] | [], _ ->
    (* we need to check whether there have been fills before,
      if so we need to adjust fill prices before getting out *)
    begin
      match fills with
      | [] -> { uncrossed_book = b; uncrossed_fills = fills; uncrossed_qty = filled_qty }
      | x::xs ->
        let fills' = x :: (adjust_fill_prices xs x.fill_price) in
      { uncrossed_book = b; uncrossed_fills = fills'; uncrossed_qty = filled_qty }
    end
  | buy::bs, sell::ss ->
    if buy.o_price >= sell.o_price then
      begin
        (* compute the fill qty and price *)
        let fill_qty = if buy.o_qty < sell.o_qty then buy.o_qty else sell.o_qty in
        let fill_price = (buy.o_price + sell.o_price) / 2 in

        (* update the orders that traded *)
        let buy' = { buy with o_qty = buy.o_qty - fill_qty } in
        let sell' = { sell with o_qty = sell.o_qty - fill_qty } in

        (* create the fills *)
        let fill1 = {
          fill_client_id = buy.o_client_id
          ; fill_qty = fill_qty
          ; fill_price = fill_price
          ; fill_order_id = buy.o_id
          ; fill_order_done = true } in

        let fill2 = {
          fill_client_id = sell.o_client_id
          ; fill_qty = fill_qty
          ; fill_price = fill_price
          ; fill_order_id = sell.o_id
          ; fill_order_done = true } in

        (* now update the books and fills *)
        let new_buys = if buy'.o_qty = 0 then bs else buy'::bs in
        let new_sells = if sell'.o_qty = 0 then ss else sell'::ss in
        let b' = {
          b_buys = new_buys
          ; b_sells = new_sells } in

        (* We should not be generating fills for implied orders - there's
          a different mechanism for that *)
        let fills' = if not buy.o_is_implied then
          fill1 :: fills else fills in
        let fills' = if not sell.o_is_implied then
          fill2 :: fills' else fills' in

        (* recursively go to the next level *)
        uncross_book b' fills' (filled_qty + fill_qty)
      end

    else
      (* nothing to do here *)
      { uncrossed_book = b; uncrossed_fills = fills; uncrossed_qty = filled_qty }
[@@measure book_measure b]
;;

let book1 = {
  b_buys = [
    (make BUY 100 55 1 (Outright OUT1) 1 false 1)
    ; (make BUY 100 50 2 (Outright OUT1) 1 false 1)
    ]
 ; b_sells = [
    (make BUY 100 54 3 (Outright OUT1) 1 false 1)
    ; (make BUY 100 54 4 (Outright OUT1) 1 false 1)
  ]
} in

uncross_book book1 [] 0

(* Returns true if the orders are sorted with respect to price *)
let rec side_price_sorted (si : side) (orders : order list) =
  match orders with
  | [] -> true
  | x::[] -> true
  | x::y::xs ->
    if si = BUY then
      begin
        if y.o_price > x.o_price && x.o_price > 0 then false else (side_price_sorted si (y::xs))
      end
    else
      begin
        if y.o_price < x.o_price && y.o_price > 0 then false else (side_price_sorted si (y::xs))
      end
;;

(* Let's make sure all the fills have this price *)
let rec fills_good_price (fills : fill list) (p : int) =
  match fills with
  | [] -> true
  | x::xs -> (x.fill_price = p) && (fills_good_price xs p)
;;

(** Let's to verify some properties *)
let fill_price_midpoint (b : book) =

  let buys_sorted = side_price_sorted BUY b.b_buys in
  let sells_sorted = side_price_sorted SELL b.b_sells in

  let result_good =
    begin
      match b.b_buys, b.b_sells with
      | [], _ -> true
      | _, [] -> true
      | x::xs, y::ys ->
        let unc_res = uncross_book b [] 0 in
        if x.o_price >= y.o_price then
          let midprice = (x.o_price + y.o_price) / 2 in
          (List.length unc_res.uncrossed_fills) > 0 && (fills_good_price unc_res.uncrossed_fills midprice)
        else
          true
    end in

  (* This is the 'punchline'... if the sides are price-sorted, then the fills will be the first midpoint *)
  (buys_sorted && sells_sorted) ==> result_good
;;


verify fill_price_midpoint

(* All no quantities get lost during uncross *)
let no_lost_qtys (b : book) =

  let rec qtys_pos_nonimp = function
    | [] -> true
    | x::xs -> x.o_qty >= 0 && not x.o_is_implied && (qtys_pos_nonimp xs) in

  let rec sum_qtys = function
    | [] -> 0
    | x::xs -> x.o_qty + (sum_qtys xs) in

  let rec sum_fill_qtys = function
    | [] -> 0
    | x::xs -> x.fill_qty + (sum_fill_qtys xs) in

  let unc_res = uncross_book b [] 0 in

  (* We need to make sure the book is non-negative *)
  let book_nonneg_nonimp = (qtys_pos_nonimp b.b_buys) && (qtys_pos_nonimp b.b_sells) in

  (* Let's sum up all of the quantities of orders before the uncross *)
  let count_before = (sum_qtys b.b_buys) + (sum_qtys b.b_sells) in

  (* And after *)
  let count_after = (sum_qtys unc_res.uncrossed_book.b_buys) +
                    (sum_qtys unc_res.uncrossed_book.b_sells) +
                    (sum_fill_qtys unc_res.uncrossed_fills) in

  book_nonneg_nonimp ==> (count_before = count_after)
;;

verify ~upto:15 no_lost_qtys

(* This is a 'side_condition' function that tells decomposition that we're only interested in cases where the
initial fills are empty *)
let cond (b : book) (fills : fill list) (filled_qty : int) =
 fills = [] && filled_qty = 0
;;

let d = Modular_decomp.top ~assuming:"cond" "uncross_book" ~prune:true [@@program];;

(* Now let's try to generate some test cases *)

(* This will auto-generate model extractor *)
Extract.eval ~quiet:true ~signature:(Event.DB.fun_id_of_str (db()) "uncross_book") ();;

#remove_doc doc_of_book;;
#remove_doc doc_of_order;;
Modular_decomp.get_regions d |> CCList.map (fun r -> r |> Modular_decomp.get_model |> Mex.of_model);;
#install_doc doc_of_order;;
#install_doc doc_of_book;;

(* Calculate somehow how big the ratio is *)
let leg_ratio (s : strategy) =
  let abs x = if x < 0 then -x else x in
  (abs s.leg1.leg_mult) + (abs s.leg2.leg_mult) + (abs s.leg3.leg_mult)
;;

(* Nearest time to expiry *)
let nearest_time_to_exp (s : strategy) =
  let exp =
    if (month_to_int (contract_expiry s.leg1.leg_sec_idx)) < (month_to_int (contract_expiry s.leg2.leg_sec_idx)) then
      contract_expiry s.leg1.leg_sec_idx
    else
      contract_expiry s.leg2.leg_sec_idx in

  if (month_to_int exp) < (month_to_int (contract_expiry s.leg2.leg_sec_idx)) then
    exp
  else
    contract_expiry s.leg2.leg_sec_idx
;;

(* Return true if s1 should implied uncross before s2 *)
let priority_strat (s1 : strategy) (s2 : strategy) =
  (*
    1. time to expiry of the nearest leg
    2. strategy types (strategies with the greater leg ratio executed first)
    3. strategy creation times *)
  if (month_to_int (nearest_time_to_exp s1)) < (month_to_int (nearest_time_to_exp s2)) then
    true
  else
    if (leg_ratio s1) > (leg_ratio s2) then
      true
    else
      s1.time_created <= s2.time_created

let transitivity s1 s2 s3 =
 ((priority_strat s1 s2) && (priority_strat s2 s3)) ==> (priority_strat s1 s3)

verify transitivity

priority_strat CX.s1 CX.s2

priority_strat CX.s2 CX.s3

priority_strat CX.s1 CX.s3

(* return the sum of volume at the highest level *)
let rec get_level_sums (orders : order list) (li : level_info option) =
  match orders with
  | [] -> li
  | x::xs ->
    begin
      match li with
      | None -> get_level_sums xs (Some {li_qty = x.o_qty; li_price = x.o_price})
      | Some l ->
        if (l.li_price = x.o_price) then
          get_level_sums xs (Some {l with li_qty = l.li_qty + x.o_qty})
        else
          li
    end

(* Return best bid/ask levels *)
let get_book_tops (b : book) =
  let bid_info = get_level_sums b.b_buys None in
  let ask_info = get_level_sums b.b_sells None in
  { bid_info; ask_info }
;;

(* Get the maximum number of strategy units here *)
(* Note that the units may have different signs, so we
 need to make sure that we have enough *)
let calc_implied_strat_order (sid : strategy_id) (s : strategy) (books : books_info) (si : side) (time : int) =
  let abs x = if x < 0 then -x else x in

  let adjust (mult : int) =
    if si = BUY then mult else -mult in

  (* *)
  let calc_max_out_mult (mult : int) (book : best_bid_ask)  =
    if mult = 0 then
      None
    else
      begin
       if (adjust mult) > 0 then
        match books.book1.bid_info with
           | Some x -> Some (x.li_qty / (abs mult))
           | None -> None
       else
           match book.ask_info with
           | Some x -> Some (x.li_qty / (abs mult))
           | None -> None
      end in

  let mult1 = calc_max_out_mult s.leg1.leg_mult books.book1 in
  let mult2 = calc_max_out_mult s.leg2.leg_mult books.book2 in
  let mult3 = calc_max_out_mult s.leg3.leg_mult books.book3 in

  (* Compute the quantity *)
  let max_strat =
    begin
       match mult1 with
       | None -> 0
       | Some x -> x
    end in
  let max_strat =
    begin
       match mult2 with
       | None -> max_strat
       | Some x -> if x < max_strat then x else max_strat
    end in
  let max_strat =
    begin
     match mult3 with
     | None -> max_strat
     | Some x -> if x < max_strat then x else max_strat
    end in

  (* Now compute the price *)
  let strat_price =
    begin
       match mult1 with
       | None -> 0
       | Some x ->
         begin
            if (adjust s.leg1.leg_mult) > 0 then
                match books.book1.bid_info with
                | Some x -> x.li_price * (adjust s.leg1.leg_mult)
                | None -> 0
            else
                match books.book1.ask_info with
                | Some x -> x.li_price * (adjust s.leg1.leg_mult)
                | None -> 0
         end
    end in
  let strat_price =
    begin
       match mult2 with
       | None -> strat_price
       | Some x ->
         begin
            if (adjust s.leg2.leg_mult) > 0 then
                match books.book2.bid_info with
                | Some x -> x.li_price * (adjust s.leg2.leg_mult) + strat_price
                | None -> strat_price
            else
                match books.book2.ask_info with
                | Some x -> x.li_price * (adjust s.leg2.leg_mult) + strat_price
                | None -> strat_price
         end
    end in
  let strat_price =
    begin
       match mult3 with
       | None -> strat_price
       | Some x ->
         begin
            if (adjust s.leg3.leg_mult) > 0 then
                match books.book3.bid_info with
                | Some x -> x.li_price * (adjust s.leg3.leg_mult) + strat_price
                | None -> strat_price
            else
                match books.book3.ask_info with
                | Some x -> x.li_price * (adjust s.leg3.leg_mult) + strat_price
                | None -> strat_price
         end
    end in

  (* Now form the new implied order here... *)
  {
   o_qty = max_strat
   ; o_price = strat_price
   ; o_id = -1
   ; o_time = time
   ; o_side = si
   ; o_client_id = -1
   ; o_inst = Strategy sid
   ; o_is_implied = true }
;;